from fastmcp import FastMCP
from dotenv import load_dotenv
import os
from pydantic import Field
from typing import Optional
import requests

load_dotenv(".env")
BDS_GRID_ENGINE_URL = os.getenv("BDS_GRID_ENGINE_URL","http://localhost:8081")

mcp = FastMCP("es-mcp",host="0.0.0.0",port=8000)

@mcp.tool()
def geometry_search(
        index_library_name:str = Field(...,min_length=1,description="数据源-一般由用户指定"),
        bit_table_names: list[str] = Field(...,min_items=1,description="es索引或别名"),
        geometry: Optional[str] = Field(...,description="geojson 参考格式：{\"type\": \"Polygon\", \"coordinates\": [[[113.85042507821815, 22.9134704440907], [113.90157492178182, 22.9134704440907], [113.90157492178182, 22.943529440473327], [113.85042507821815, 22.943529440473327], [113.85042507821815, 22.9134704440907]]]}"),
        page_num: Optional[int] = Field(...,description="页码 从1开始"),
        page_size: Optional[int] = Field(...,description="每页条数"),
    ) -> list[dict]:
    """
    空间数据检索
    index_library_name: 数据源-一般由用户指定
    bit_table_names: es索引或别名
    geometry: geojson 参考格式：{"type": "Polygon", "coordinates": [[[113.85042507821815, 22.9134704440907], [113.90157492178182, 22.9134704440907], [113.90157492178182, 22.943529440473327], [113.85042507821815, 22.943529440473327], [113.85042507821815, 22.9134704440907]]]}
    page_num: 页码 从1开始
    page_size: 每页条数
    """
    router = '/satxspace-gdm2/data/search/geometry'
    url = f"{BDS_GRID_ENGINE_URL}{router}"
    params = {
        "index_library_name": index_library_name,
        "big_table_names": ",".join(bit_table_names),
        "geometry": geometry,
        "page_num": page_num,
        "page_size": page_size,
    }

    response = requests.post(url, params=params)
    return response.json()

@mcp.tool()
def get_data_source_list() -> list[dict]:
    """
    获取数据源列表 - 即index_library_name
    """
    router = '/satxspace-gdm2/nodemanage/find'
    url = f"{BDS_GRID_ENGINE_URL}{router}"
    params = {
        "index_library_name": "",
    }
    response = requests.post(url, params=params)
    return response.json()
    
@mcp.tool()
def get_layer_by_indexLibraryName(index_library_name:str) -> list[dict]:
    """
    获取一个数据源中的所有图层
    """
    router = '/satxspace-gdm2/data/search/geometry/whichLayer'
    url = f"{BDS_GRID_ENGINE_URL}{router}"
    params = {
        "index_library_name": index_library_name,
    }
    response = requests.post(url, params=params)
    return response.json()


    
if __name__ == "__main__":
    mcp.run(transport="sse")